#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include "mpi.h"

void print_matrix(double *matrix, int n) {
    for (int i = 0; i < n; i++) {
        for (int j = 0; j < n; j++) {
            printf("%f ", matrix[i * n + j]);
        }
        printf("\n");
    }
}

int main(int argc, char *argv[]) {
    int numprocs, myid;
    MPI_Status status;

    MPI_Init(&argc, &argv);
    MPI_Comm_rank(MPI_COMM_WORLD, &myid);
    MPI_Comm_size(MPI_COMM_WORLD, &numprocs);

    int N;
    double *A, *B, *C;
    int rows_per_proc;

    if (myid == 0) {
        printf("Enter the matrix size N: ");
        scanf("%d", &N);

        A = (double *)malloc(N * N * sizeof(double));
        B = (double *)malloc(N * N * sizeof(double));
        C = (double *)malloc(N * N * sizeof(double));

        printf("Enter matrix A:\n");
        for (int i = 0; i < N * N; i++) {
            scanf("%lf,", &A[i]);
        }

        printf("Enter matrix B:\n");
        for (int i = 0; i < N * N; i++) {
            scanf("%lf,", &B[i]);
        }

        printf("Matrix A:\n");
        print_matrix(A, N);
        printf("Matrix B:\n");
        print_matrix(B, N);

        rows_per_proc = N / numprocs;

        // 发送矩阵大小 N 和矩阵 B 给所有进程
        for (int dest = 1; dest < numprocs; dest++) {
            MPI_Send(&N, 1, MPI_INT, dest, 0, MPI_COMM_WORLD);
            MPI_Send(B, N * N, MPI_DOUBLE, dest, 1, MPI_COMM_WORLD);
        }
    } else {
        // 接收矩阵大小 N 和矩阵 B
        MPI_Recv(&N, 1, MPI_INT, 0, 0, MPI_COMM_WORLD, &status);
        B = (double *)malloc(N * N * sizeof(double));
        MPI_Recv(B, N * N, MPI_DOUBLE, 0, 1, MPI_COMM_WORLD, &status);
    }

    rows_per_proc = N / numprocs;
    double *local_A = (double *)malloc(rows_per_proc * N * sizeof(double));
    double *local_C = (double *)malloc(rows_per_proc * N * sizeof(double));

    if (myid == 0) {
        // 发送矩阵 A 的部分给每个进程
        for (int dest = 1; dest < numprocs; dest++) {
            MPI_Send(A + dest * rows_per_proc * N, rows_per_proc * N, MPI_DOUBLE, dest, 2, MPI_COMM_WORLD);
        }

        memcpy(local_A, A, rows_per_proc * N * sizeof(double));
    } else {
        MPI_Recv(local_A, rows_per_proc * N, MPI_DOUBLE, 0, 2, MPI_COMM_WORLD, &status);
    }

    for (int i = 0; i < rows_per_proc; i++) {
        for (int j = 0; j < N; j++) {
            local_C[i * N + j] = 0.0;
            for (int k = 0; k < N; k++) {
                local_C[i * N + j] += local_A[i * N + k] * B[k * N + j];
            }
        }
    }

    if (myid == 0) {
        memcpy(C, local_C, rows_per_proc * N * sizeof(double));
        for (int source = 1; source < numprocs; source++) {
            MPI_Recv(C + source * rows_per_proc * N, rows_per_proc * N, MPI_DOUBLE, source, 3, MPI_COMM_WORLD, &status);
        }

        printf("Matrix C:\n");
        print_matrix(C, N);

        free(A);
        free(B);
        free(C);
    } else {
        MPI_Send(local_C, rows_per_proc * N, MPI_DOUBLE, 0, 3, MPI_COMM_WORLD);
    }

    free(local_A);
    free(local_C);

    MPI_Finalize();
    return 0;
}
